import os
import pandas as pd
from openpyxl import load_workbook
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import warnings

warnings.filterwarnings('ignore')


class ExcelColumnSelectorWithTable:
    def __init__(self, root):
        self.root = root
        self.root.title("Excel列名选择工具 - 表格预览版")
        self.root.geometry("1200x800")

        # 变量初始化
        self.file_path = None
        self.sheet_name = None
        self.df = None
        self.header_row = 0  # 默认第一行为列名行
        self.selected_columns = []
        self.output_file = "统计报告.xlsx"
        self.report_data = []

        # 创建界面
        self.create_widgets()

    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding=10)
        main_frame.pack(fill="both", expand=True)

        # 左侧文件处理区域
        left_frame = ttk.LabelFrame(main_frame, text="文件处理", padding=10)
        left_frame.pack(side="left", fill="y", padx=5, pady=5)

        # 文件选择区域
        file_frame = ttk.Frame(left_frame)
        file_frame.pack(fill="x", pady=5)

        ttk.Button(file_frame, text="选择Excel文件", command=self.select_file).pack(side="left", padx=5)
        self.file_label = ttk.Label(file_frame, text="未选择文件")
        self.file_label.pack(side="left", padx=5)

        # Sheet选择区域
        self.sheet_frame = ttk.Frame(left_frame)
        self.sheet_frame.pack(fill="x", pady=5)
        ttk.Label(self.sheet_frame, text="选择Sheet:").pack(side="left", padx=5)
        self.sheet_combo = ttk.Combobox(self.sheet_frame, state="readonly", width=20)
        self.sheet_combo.pack(side="left", padx=5)
        self.sheet_combo.bind("<<ComboboxSelected>>", self.on_sheet_select)

        # 列名行选择区域
        row_frame = ttk.Frame(left_frame)
        row_frame.pack(fill="x", pady=5)
        ttk.Label(row_frame, text="列名所在行:").pack(side="left", padx=5)
        self.row_spinbox = ttk.Spinbox(row_frame, from_=1, to=20, width=5, command=self.on_row_change)
        self.row_spinbox.set(1)
        self.row_spinbox.pack(side="left", padx=5)
        ttk.Button(row_frame, text="应用", command=self.apply_row_change).pack(side="left", padx=5)

        # 数据预览区域
        preview_frame = ttk.LabelFrame(left_frame, text="数据预览 (前10行)", padding=5)
        preview_frame.pack(fill="both", expand=True, pady=5)

        # 创建表格预览
        self.preview_tree = ttk.Treeview(preview_frame, height=10)
        preview_vscroll = ttk.Scrollbar(preview_frame, orient="vertical", command=self.preview_tree.yview)
        preview_hscroll = ttk.Scrollbar(preview_frame, orient="horizontal", command=self.preview_tree.xview)
        self.preview_tree.configure(yscrollcommand=preview_vscroll.set, xscrollcommand=preview_hscroll.set)

        self.preview_tree.pack(side="left", fill="both", expand=True)
        preview_vscroll.pack(side="right", fill="y")
        preview_hscroll.pack(side="bottom", fill="x")

        # 列选择区域
        col_frame = ttk.LabelFrame(left_frame, text="选择要合并的列", padding=5)
        col_frame.pack(fill="both", expand=True, pady=5)

        self.col_listbox = tk.Listbox(col_frame, selectmode=tk.MULTIPLE, height=8)
        col_scrollbar = ttk.Scrollbar(col_frame, orient="vertical", command=self.col_listbox.yview)
        self.col_listbox.configure(yscrollcommand=col_scrollbar.set)

        self.col_listbox.pack(side="left", fill="both", expand=True)
        col_scrollbar.pack(side="right", fill="y")

        # 按钮区域
        btn_frame = ttk.Frame(left_frame)
        btn_frame.pack(fill="x", pady=5)

        ttk.Button(btn_frame, text="全选", command=self.select_all).pack(side="left", padx=5)
        ttk.Button(btn_frame, text="全不选", command=self.deselect_all).pack(side="left", padx=5)
        ttk.Button(btn_frame, text="添加到报告", command=self.add_to_report).pack(side="right", padx=5)

        # 右侧报告区域
        right_frame = ttk.LabelFrame(main_frame, text="统计报告", padding=10)
        right_frame.pack(side="right", fill="both", expand=True, padx=5, pady=5)

        # 创建表格显示报告
        columns = ("文件名称", "所有字段", "字段数量")
        self.report_tree = ttk.Treeview(right_frame, columns=columns, show="headings", height=15)

        for col in columns:
            self.report_tree.heading(col, text=col)
            self.report_tree.column(col, width=150)

        # 添加滚动条
        scrollbar_report = ttk.Scrollbar(right_frame, orient="vertical", command=self.report_tree.yview)
        self.report_tree.configure(yscrollcommand=scrollbar_report.set)

        self.report_tree.pack(side="left", fill="both", expand=True)
        scrollbar_report.pack(side="right", fill="y")

        # 底部按钮区域
        bottom_frame = ttk.Frame(self.root)
        bottom_frame.pack(fill="x", pady=10)

        ttk.Button(bottom_frame, text="保存报告", command=self.save_report).pack(side="right", padx=10)
        ttk.Button(bottom_frame, text="清空报告", command=self.clear_report).pack(side="right", padx=10)
        ttk.Button(bottom_frame, text="批量处理文件夹", command=self.batch_process).pack(side="left", padx=10)

    def select_file(self):
        file_path = filedialog.askopenfilename(
            title="选择Excel文件",
            filetypes=[("Excel文件", "*.xlsx *.xls"), ("所有文件", "*.*")]
        )

        if file_path:
            self.file_path = file_path
            self.file_label.config(text=os.path.basename(file_path))
            self.load_file_info()

    def load_file_info(self):
        try:
            # 检查是否为Word文件
            if self.file_path.lower().endswith(('.doc', '.docx')):
                self.add_to_report_with_message("Word文件跳过")
                return

            # 获取所有sheet信息
            wb = load_workbook(self.file_path, read_only=True)
            sheet_names = wb.sheetnames
            wb.close()

            # 更新sheet选择框
            self.sheet_combo['values'] = sheet_names
            if len(sheet_names) > 0:
                self.sheet_combo.set(sheet_names[0])
                self.sheet_name = sheet_names[0]
                self.load_sheet_data()
        except Exception as e:
            messagebox.showerror("错误", f"读取文件时出错: {str(e)}")

    def on_sheet_select(self, event):
        self.sheet_name = self.sheet_combo.get()
        self.load_sheet_data()

    def on_row_change(self):
        # 当行号改变时，更新预览
        try:
            row_num = int(self.row_spinbox.get()) - 1  # 转换为0-based索引
            if self.df is not None and 0 <= row_num < len(self.df):
                self.update_preview(row_num)
        except ValueError:
            pass

    def apply_row_change(self):
        # 应用行号更改，重新加载数据
        try:
            self.header_row = int(self.row_spinbox.get()) - 1  # 转换为0-based索引
            self.load_sheet_data()
        except ValueError:
            messagebox.showerror("错误", "请输入有效的行号")

    def load_sheet_data(self):
        try:
            # 读取Excel文件，尝试使用指定的行作为列名
            self.df = pd.read_excel(
                self.file_path,
                sheet_name=self.sheet_name,
                header=self.header_row
            )

            # 更新列选择列表
            self.update_column_list()

            # 更新数据预览
            self.update_preview(self.header_row)

        except Exception as e:
            messagebox.showerror("错误", f"读取数据时出错: {str(e)}")

    def update_column_list(self):
        # 更新列选择列表框
        self.col_listbox.delete(0, tk.END)
        if self.df is not None:
            for col in self.df.columns:
                self.col_listbox.insert(tk.END, col)

            # 默认选择所有列
            self.select_all()

    def update_preview(self, header_row):
        # 清空预览表格
        for item in self.preview_tree.get_children():
            self.preview_tree.delete(item)

        # 清空列定义
        self.preview_tree['columns'] = []

        if self.df is not None:
            # 显示前10行数据
            preview_df = self.df.head(10)

            # 设置表格列
            columns = list(preview_df.columns)
            self.preview_tree['columns'] = columns

            # 设置列标题和宽度
            for col in columns:
                self.preview_tree.heading(col, text=col)
                self.preview_tree.column(col, width=100, minwidth=50)

            # 添加数据行
            for i, row in preview_df.iterrows():
                # 处理NaN值，显示为空字符串
                values = [str(row[col]) if pd.notna(row[col]) else "" for col in columns]
                self.preview_tree.insert("", "end", values=values)

    def select_all(self):
        self.col_listbox.selection_set(0, tk.END)

    def deselect_all(self):
        self.col_listbox.selection_clear(0, tk.END)

    def get_selected_columns(self):
        selected_indices = self.col_listbox.curselection()
        return [self.df.columns[i] for i in selected_indices] if self.df is not None else []

    def add_to_report(self):
        if not self.file_path:
            messagebox.showwarning("警告", "请先选择Excel文件")
            return

        if not self.sheet_name:
            messagebox.showwarning("警告", "请先选择Sheet")
            return

        selected_columns = self.get_selected_columns()
        if not selected_columns:
            messagebox.showwarning("警告", "请至少选择一列")
            return

        # 生成字段字符串和数量
        columns_str = "、".join(selected_columns)
        column_count = len(selected_columns)

        # 添加到报告
        self.report_tree.insert("", "end", values=(
            os.path.basename(self.file_path),
            columns_str,
            column_count
        ))

        # 保存到数据列表
        self.report_data.append({
            '文件名称': os.path.basename(self.file_path),
            '所有字段': columns_str,
            '字段数量': column_count
        })

    def add_to_report_with_message(self, message):
        self.report_tree.insert("", "end", values=(
            os.path.basename(self.file_path),
            message,
            0
        ))

        # 保存到数据列表
        self.report_data.append({
            '文件名称': os.path.basename(self.file_path),
            '所有字段': message,
            '字段数量': 0
        })

    def clear_report(self):
        # 清空报告
        for item in self.report_tree.get_children():
            self.report_tree.delete(item)
        self.report_data = []

    def save_report(self):
        if not self.report_data:
            messagebox.showwarning("警告", "报告为空")
            return

        # 选择保存位置
        save_path = filedialog.asksaveasfilename(
            title="保存报告",
            defaultextension=".xlsx",
            filetypes=[("Excel文件", "*.xlsx"), ("所有文件", "*.*")]
        )

        if save_path:
            try:
                # 创建DataFrame并保存
                df = pd.DataFrame(self.report_data)
                df.to_excel(save_path, index=False)
                messagebox.showinfo("成功", f"报告已保存到: {save_path}")
            except Exception as e:
                messagebox.showerror("错误", f"保存报告时出错: {str(e)}")

    def batch_process(self):
        # 批量处理文件夹中的文件
        folder_path = filedialog.askdirectory(title="选择包含Excel文件的文件夹")

        if not folder_path:
            return

        # 获取文件夹中的所有文件
        files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
        excel_files = [f for f in files if f.lower().endswith(('.xlsx', '.xls'))]
        word_files = [f for f in files if f.lower().endswith(('.doc', '.docx'))]

        # 处理Excel文件
        for file in excel_files:
            file_path = os.path.join(folder_path, file)
            self.file_path = file_path
            self.file_label.config(text=os.path.basename(file_path))

            try:
                # 获取所有sheet信息
                wb = load_workbook(file_path, read_only=True)
                sheet_names = wb.sheetnames
                wb.close()

                # 检查是否有多个有内容的sheet
                sheets_with_content = []
                for sheet in sheet_names:
                    try:
                        df_check = pd.read_excel(file_path, sheet_name=sheet, nrows=5)
                        if not df_check.empty and not df_check.isnull().all().all():
                            sheets_with_content.append(sheet)
                    except:
                        continue

                # 如果有多个有内容的sheet，跳过
                if len(sheets_with_content) > 1:
                    self.add_to_report_with_message("多sheet文件跳过")
                    continue

                # 处理第一个有内容的sheet
                if sheets_with_content:
                    self.sheet_name = sheets_with_content[0]
                    self.load_sheet_data()

                    # 自动选择所有列
                    self.select_all()

                    # 添加到报告
                    self.add_to_report()
                else:
                    self.add_to_report_with_message("无有效内容的sheet")

            except Exception as e:
                self.add_to_report_with_message(f"处理文件时出错: {str(e)}")

        # 处理Word文件
        for file in word_files:
            self.add_to_report_with_message("Word文件跳过")

        messagebox.showinfo("完成", f"已处理 {len(excel_files)} 个Excel文件和 {len(word_files)} 个Word文件")


def main():
    root = tk.Tk()
    app = ExcelColumnSelectorWithTable(root)
    root.mainloop()


if __name__ == "__main__":
    main()